'''
实现低分辨率和高分辨率像元匹配
'''
from osgeo import gdal, osr
import numpy as np


# 读取tif文件
def read_tif(file_path):
    dataset = gdal.Open(file_path)
    data = dataset.ReadAsArray()
    geotransform = dataset.GetGeoTransform()
    return data, geotransform


# 写入tif文件
def write_tif(file_path, data, geotransform, nodata):
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    dataset = driver.Create(file_path, cols, rows, 1, gdal.GDT_Float32)
    dataset.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    dataset.SetProjection(prj.ExportToWkt())
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    del dataset
    print("写入成功：{}".format(file_path))


# 根据像元行列号计算经纬度
def pixel2coord(geo_transform, x, y):
    x_origin, pixel_width, x_rotation, y_origin, y_rotation, pixel_height = geo_transform
    x_geo = x_origin + pixel_width * x
    y_geo = y_origin + pixel_height * y
    return x_geo, y_geo


# 根据经纬度计算像元行列号
def coord2pixel(geo_transform, x_geo, y_geo):
    x = int((x_geo - geo_transform[0]) / geo_transform[1])
    y = int((y_geo - geo_transform[3]) / geo_transform[5])
    return x, y


# 获取低分辨率像元对应所有高分辨率像元值
def get_high_res_pixels(high_res_geo_transform, high_res_data, low_res_geo_transform, x_low_res_pixel, y_low_res_pixel, low_res_to_high_res_ratio):
    # 计算低分辨率影像像元在高分辨率影像中的范围
    x_low_res_geo, y_low_res_geo = pixel2coord(low_res_geo_transform, x_low_res_pixel, y_low_res_pixel)
    x_low_res_geo = x_low_res_geo + low_res_geo_transform[1] / 2  # 左右偏移
    y_low_res_geo = y_low_res_geo + low_res_geo_transform[5] / 2   # 上下偏移

    # 计算行列号范围
    # high_res_inv_geo_transform = gdal.InvGeoTransform(high_res_geo_transform)
    # high_res_x, high_res_y = gdal.ApplyGeoTransform(high_res_inv_geo_transform, x_low_res_geo, y_low_res_geo)
    high_res_x, high_res_y = coord2pixel(high_res_geo_transform, x_low_res_geo, y_low_res_geo)
    # if 0 <= high_res_x <= high_res_data.shape[1] and 0 <= high_res_y <= high_res_data.shape[0]:
    #     print(x_low_res_pixel, y_low_res_pixel, high_res_x, high_res_y)
    high_res_x_min = int(high_res_x - low_res_to_high_res_ratio / 2)
    high_res_x_max = int(high_res_x + low_res_to_high_res_ratio / 2)
    high_res_y_min = int(high_res_y - low_res_to_high_res_ratio / 2)
    high_res_y_max = int(high_res_y + low_res_to_high_res_ratio / 2)

    # 排除边缘像元的影响
    if high_res_x_min < 0 and high_res_x_max > 0:
        high_res_x_min = 0
    if high_res_x_min < high_res_data.shape[1] < high_res_x_max:
        high_res_x_max = high_res_data.shape[1]
    if high_res_y_min < 0 and high_res_y_max > 0:
        high_res_y_min = 0
    if high_res_y_min < high_res_data.shape[0] < high_res_y_max:
        high_res_y_max = high_res_data.shape[0]

    # 获取范围内的像元值
    if 0 <= high_res_x_min <= high_res_data.shape[1] and 0 <= high_res_y_min <= high_res_data.shape[0]:
        # print(high_res_x_min, high_res_y_min, high_res_x_max, high_res_y_max)
        high_res_pixels = high_res_data[high_res_y_min:high_res_y_max, high_res_x_min:high_res_x_max]
        return high_res_pixels
    else:
        return []


# 主函数
def main(CCI_image, GLASS_image, output_cci_measure_1km, output_glass_average_1km):
    # 读取高分辨率和低分辨率的影像数据
    high_res_data, high_res_geotrans = read_tif(GLASS_image)
    low_res_data, low_res_geotrans = read_tif(CCI_image)
    print(high_res_geotrans)
    print(low_res_geotrans)

    low_res_to_high_res_ratio = int(low_res_geotrans[1] / high_res_geotrans[1])
    print("像元比例：{}".format(low_res_to_high_res_ratio))
    # 计算25km分辨率下对应GLASS数据的平均值
    glass_average_25km = np.zeros(low_res_data.shape)
    for y in range(0, low_res_data.shape[0]):
        for x in range(0, low_res_data.shape[1]):
            high_res_pixels = get_high_res_pixels(high_res_geotrans, high_res_data, low_res_geotrans, x, y, low_res_to_high_res_ratio)
            if len(high_res_pixels) != 0:
                glass_average_25km[y, x] = np.mean(high_res_pixels)

    glass_average_1km = np.zeros(high_res_data.shape)
    # 计算1km分辨率下对应CCI数据
    cci_measure_1km = np.zeros(high_res_data.shape)
    for y in range(0, high_res_data.shape[0]):
        for x in range(0, high_res_data.shape[1]):
            # print(high_res_data[y, x])
            # 只读取glass有值的对应值
            if high_res_data[y, x] != -1:
                # 计算glass像元对应的经纬度
                x_glass_geo, y_glass_geo = pixel2coord(high_res_geotrans, x, y)
                # 计算该经纬度对应的cci像元行列号
                x_cci, y_cci = coord2pixel(low_res_geotrans, x_glass_geo, y_glass_geo)
                # 获取该点对应的cci数据
                cci_measure_1km[y, x] = low_res_data[y_cci, x_cci]
                glass_average_1km[y, x] = glass_average_25km[y_cci, x_cci]

    # 去除异常值
    cci_measure_1km[cci_measure_1km == 0] = -9999
    glass_average_1km[glass_average_1km == 0] = -1
    # 生成输出的1km cci映射影像
    write_tif(output_cci_measure_1km + ".tif", cci_measure_1km, high_res_geotrans, -9999)
    # 生成输出的1km glass平均值影像
    write_tif(output_glass_average_1km + ".tif", glass_average_1km, high_res_geotrans, -1)
    print("映射完成！")


if __name__ == "__main__":
    # main()
    CCI_path = r"ESACCI-SOILMOISTURE-L3S-SSMV-COMBINED-20130728000000-fv04.2.tif"
    GLASS_path = r"GLASS02A06.V40.A2013209.h25v05.2017101.tif"
    output_cci_measure_1km = "cci_measure_1km"
    output_glass_average_1km = "glass_average_1km"
    main(CCI_path, GLASS_path, output_cci_measure_1km, output_glass_average_1km)
